package com.fliad.ai.core;

import cn.hutool.core.util.IdUtil;
import com.fliad.ai.dto.taskmodel.NodeReport;
import com.fliad.ai.dto.taskmodel.Snapshot;
import com.fliad.ai.dto.taskmodel.TaskReportOutput;
import org.noear.snack.ONode;
import org.noear.solon.expression.snel.SnEL;
import org.noear.solon.flow.*;
import org.noear.solon.flow.stateful.StateController;
import org.noear.solon.flow.stateful.StateRepository;
import org.noear.solon.flow.stateful.driver.StatefulSimpleFlowDriver;

import java.util.Map;

/**
 * ai 流程驱动器定制
 */
public class AiFlowDriver extends StatefulSimpleFlowDriver {

    public AiFlowDriver(StateRepository stateRepository, StateController stateController, Evaluation evaluation, Container container) {
        super(stateRepository, stateController, evaluation, container);
    }


    @Override
    public void onNodeStart(FlowContext context, Node node) {
        // 节点开始增加任务快照
        TaskReportOutput report = context.get("report");
        NodeReport nodeReport = report.addNodeReport(new NodeReport(node.getId()));
        nodeReport.addSnapshot(new Snapshot(IdUtil.getSnowflakeNextIdStr(), node.getId()));
        if (node.getType() == NodeType.START) {
            nodeReport.getLastSnapshot().setInputs(report.getInputs());
            nodeReport.setTerminated(true);
        } else if (node.getType() == NodeType.INCLUSIVE) {
            nodeReport.getLastSnapshot().setInputs(report.getInputs());
            nodeReport.setTerminated(true);
        }
    }


    @Override
    public void onNodeEnd(FlowContext context, Node node) {
        // 节点结束调整快照状态
        TaskReportOutput report = context.get("report");
        Map<String, Object> variableContext = context.get("variableContext");
        NodeReport nodeReport = report.getReports().get(node.getId());
        nodeReport.setTerminated(true);
        // 结束工作流
        if (node.getType() == NodeType.END) {
            ONode inputsValues = (ONode) node.getMetas().get("inputsValues");
            if (!inputsValues.isNull()) {
                inputsValues.forEach((key, value) -> {
                    // 引用变量
                    if ("ref".equals(value.get("type").getString())) {
                        nodeReport.getLastSnapshot().getOutputs().put(key,
                                SnEL.eval(String.format("%s.%s", value.get("content").get(0).getString(),
                                        value.get("content").get(1).getString()), variableContext)
                        );
                        report.setOutputs(nodeReport.getLastSnapshot().getOutputs());
                    }
                });
            }

            report.getWorkflowStatus().setTerminated(true);
        }
    }

    public static AiFlowDriver.Builder builder1() {
        return new AiFlowDriver.Builder();
    }

    public static class Builder {
        private StateRepository stateRepository;
        private StateController stateController;
        private Evaluation evaluation;
        private Container container;

        /**
         * 设置状态仓库
         */
        public AiFlowDriver.Builder stateRepository(StateRepository stateRepository) {
            this.stateRepository = stateRepository;
            return this;
        }

        /**
         * 设置状态控制器
         */
        public AiFlowDriver.Builder stateController(StateController stateController) {
            this.stateController = stateController;
            return this;
        }

        /**
         * 设置评估器
         */
        public AiFlowDriver.Builder evaluation(Evaluation evaluation) {
            this.evaluation = evaluation;
            return this;
        }

        /**
         * 设置容器
         */
        public AiFlowDriver.Builder container(Container container) {
            this.container = container;
            return this;
        }

        /**
         * 构建
         */
        public AiFlowDriver build() {
            return new AiFlowDriver(
                    stateRepository,
                    stateController,
                    evaluation,
                    container);
        }
    }

}
